library(dplyr)
library(ggplot2)
library(ggrepel)
library(stringr)

source("HNC_metadata_tidy.R")

if (!file.exists("output")) {
  dir.create("output")
}

data <- data %>%
  # Change labels
  mutate(
    subsite = factor(subsite, labels = c("OC", "OPC", "HPX", "LYX")),
    hpv_pos_all = ifelse(subsite == "OPC", hpv_pos_opc, "Not applicable"),
    hpv_pos_all = factor(hpv_pos_all, levels = c("Negative", "Positive", "Not applicable"))
  )

# Define variables for plotting
var <- c("tobacco", "alcohol_ever", "tob_alc", "hpv_pos_all", "oral_cat", "bmi_cat", "hotdrinks")

for (v in var) {
  prop <- data %>%
    group_by(var = get(v)) %>%
    tally() %>%
    na.omit()
  prop$fraction <- prop$n / sum(prop$n)
  prop$ymax <- cumsum(prop$fraction)
  prop$ymin <- c(0, head(prop$ymax, n = -1))
  prop$labelPosition <- (prop$ymax + prop$ymin) / 2
  prop$label <- paste0(prop$var, "\n n = ", prop$n)

  palette <- scales::brewer_pal(palette = 4)(length(levels(data[, v])))
  names(palette) <- levels(data[, v])
  palette["Not applicable"] <- "#E0E0E0"

  # Plot for overall dataset
  p <- ggplot(prop, aes(ymax = ymax, ymin = ymin, xmax = 4, xmin = 3, fill = var)) +
    geom_rect() +
    geom_text_repel(
      x = 3.5, aes(
        y = labelPosition,
        label = str_to_sentence(str_replace_all(label, "_", " "))
      ),
      size = 4
    ) +
    scale_fill_manual(values = palette) +
    coord_polar(theta = "y") +
    xlim(c(2, 4)) +
    labs(title = str_replace_all(v, "_", " "), fill = v) +
    theme_void() +
    theme(
      plot.title = element_text(face = "bold", size = 16, hjust = 0.5),
      legend.position = "none"
    )

  plots <- NULL
  for (s in unique(data$subsite)) {
    prop <- data %>%
      filter(subsite == s) %>%
      group_by(var = get(v)) %>%
      tally() %>%
      na.omit()
    prop$fraction <- prop$n / sum(prop$n)
    prop$ymax <- cumsum(prop$fraction)
    prop$ymin <- c(0, head(prop$ymax, n = -1))
    prop$labelPosition <- (prop$ymax + prop$ymin) / 2
    prop$label <- paste0(prop$var, "\n n = ", prop$n)

    # Plot by anatomical subsite
    p_s <- ggplot(prop, aes(ymax = ymax, ymin = ymin, xmax = 4, xmin = 3, fill = var)) +
      geom_rect() +
      geom_text(x = 2, y = 1, label = s, size = 4) +
      scale_fill_manual(values = palette) +
      coord_polar(theta = "y") +
      labs(fill = v) +
      xlim(c(2, 4)) +
      theme_void() +
      theme(legend.position = "None")
    plots[[s]] <- p_s
  }
  plot_s <- ggarrange(plotlist = plots, ncol = 2, nrow = 2)
  plot_all <- ggarrange(p, plot_s, nrow = 1, widths = c(1.4, 1), align = "h")
  plot(plot_all)
  ggsave(paste0("output/MUT_EAC_", v, "_", Sys.Date(), ".pdf"), device = "pdf", width = 10, height = 5)
}
